#!/bin/bash

# Qwen2.5-32B 基础环境部署脚本 (适用于 Ubuntu 22.04 + A100 80G)
set -e  # 遇到错误立即退出

echo "开始安装Qwen2.5-32B基础环境..."

# 1. 更新系统并安装基础依赖
echo "步骤1/6: 更新系统包及安装基础工具..."
sudo apt update && sudo apt upgrade -y
sudo apt install -y git-lfs wget curl build-essential ca-certificates

# 2. 安装Miniconda (Python环境管理)
echo "步骤2/6: 安装Miniconda..."
if ! command -v conda &> /dev/null; then
    wget https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-x86_64.sh -O miniconda.sh
    bash miniconda.sh -b -p $HOME/miniconda
    rm miniconda.sh

    # === 关键修改：立即将Miniconda的路径添加到当前Shell的PATH中 ===
    export PATH="$HOME/miniconda/bin:$PATH"
    # 同时也写入.bashrc以供将来使用
    echo 'export PATH="$HOME/miniconda/bin:$PATH"' >> ~/.bashrc

    # 初始化 conda（支持 conda activate）
    echo "初始化 Conda..."
    conda init bash
    # 对于 zsh 用户也初始化
    if [ -n "$ZSH_VERSION" ]; then
        conda init zsh
    fi
    
    echo "注意: 如果当前 shell 无法使用 'conda activate'，请运行:"
    echo "  source ~/.bashrc  (对于 bash)"
    echo "  或"
    echo "  source ~/.zshrc   (对于 zsh)"
    echo "  或者重新打开终端"

else
    echo "Conda 已安装，跳过。"
    # 即使 conda 已安装，也检查是否已初始化
    if ! grep -q "conda initialize" ~/.bashrc 2>/dev/null; then
        echo "检测到 Conda 未初始化，正在初始化..."
        conda init bash
        if [ -n "$ZSH_VERSION" ]; then
            conda init zsh
        fi
        echo "请运行: source ~/.bashrc 或重新打开终端"
    fi
fi

# 3. 创建并激活Conda虚拟环境 (Python 3.10)
echo "步骤3/6: 创建Conda虚拟环境..."
# 现在conda命令在当前脚本中应该可用了
conda create -n qwen32b python=3.10 -y

# 4. 配置PyPI镜像加速及升级pip
echo "步骤4/6: 配置Python环境..."
pip config set global.index-url https://pypi.tuna.tsinghua.edu.cn/simple
python -m pip install --upgrade pip

# 5. 安装PyTorch (CUDA 12.1) 及核心库
echo "步骤5/6: 安装PyTorch和vLLM..."
pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu121
pip install vllm

# 6. 下载Qwen2.5-32B模型 (通过ModelScope)
echo "步骤6/6: 下载Qwen2.5-32B模型..."
pip install modelscope
python -c "
from modelscope import snapshot_download
model_dir = snapshot_download('Qwen/Qwen2.5-32B', cache_dir='./models')
print(f'模型已下载至: {model_dir}')
"

echo "基础环境部署完成！"
echo ""
echo "========================================="
echo "下一步：激活环境"
echo "========================================="
echo "如果遇到 'conda activate' 错误，请运行修复脚本："
echo "  bash $(dirname "$0")/fix_conda_init.sh"
echo ""
echo "或者手动运行："
echo "  source ~/.bashrc  (bash 用户)"
echo "  或"
echo "  source ~/.zshrc   (zsh 用户)"
echo "  或者重新打开终端"
echo ""
echo "然后激活环境："
echo "  conda activate qwen32b"
echo ""
echo "========================================="
echo "测试环境"
echo "========================================="
echo "1. 运行测试脚本验证安装："
echo "   bash $(dirname "$0")/test_qwen32b.sh"
echo "   或"
echo "   python $(dirname "$0")/quick_test_qwen.py"
echo ""
echo "2. 启动 vLLM API 服务（可选）："
echo "   python -m vllm.entrypoints.openai.api_server \\"
echo "     --model ./models/Qwen/Qwen2.5-32B \\"
echo "     --trust-remote-code"
echo "========================================="

